{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/retrievers/reciprocal_rerank_fusion.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reciprocal Rerank Fusion Retriever\n",
    "\n",
    "In this example, we walk through how you can combine retrieval results from multiple queries and multiple indexes. \n",
    "\n",
    "The retrieved nodes will be reranked according to the `Reciprocal Rerank Fusion` algorithm demonstrated in this [paper](https://plg.uwaterloo.ca/~gvcormac/cormacksigir09-rrf.pdf). It provides an effecient method for rerranking retrieval results without excessive computation or reliance on external models.\n",
    "\n",
    "Full credits go to @Raduaschl on github for their [example implementation here](https://github.com/Raudaschl/rag-fusion)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai\n",
    "%pip install llama-index-retrievers-bm25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import openai\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-02-12 17:59:58--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8003::154, 2606:50c0:8001::154, 2606:50c0:8002::154, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8003::154|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 75042 (73K) [text/plain]\n",
      "Saving to: ‘data/paul_graham/paul_graham_essay.txt’\n",
      "\n",
      "data/paul_graham/pa 100%[===================>]  73.28K   327KB/s    in 0.2s    \n",
      "\n",
      "2024-02-12 17:59:59 (327 KB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we will setup a vector index over the documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "\n",
    "splitter = SentenceSplitter(chunk_size=256)\n",
    "\n",
    "index = VectorStoreIndex.from_documents(documents, transformations=[splitter])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Hybrid Fusion Retriever\n",
    "\n",
    "In this step, we fuse our index with a BM25 based retriever. This will enable us to capture both semantic relations and keywords in our input queries.\n",
    "\n",
    "Since both of these retrievers calculate a score, we can use the reciprocal rerank algorithm to re-sort our nodes without using an additional models or excessive computation.\n",
    "\n",
    "This setup will also query 4 times, once with your original query, and generate 3 more queries.\n",
    "\n",
    "By default, it uses the following prompt to generate extra queries:\n",
    "\n",
    "```python\n",
    "QUERY_GEN_PROMPT = (\n",
    "    \"You are a helpful assistant that generates multiple search queries based on a \"\n",
    "    \"single input query. Generate {num_queries} search queries, one on each line, \"\n",
    "    \"related to the following input query:\\n\"\n",
    "    \"Query: {query}\\n\"\n",
    "    \"Queries:\\n\"\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we create our retrievers. Each will retrieve the top-2 most similar nodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.retrievers.bm25 import BM25Retriever\n",
    "\n",
    "vector_retriever = index.as_retriever(similarity_top_k=2)\n",
    "\n",
    "bm25_retriever = BM25Retriever.from_defaults(\n",
    "    docstore=index.docstore, similarity_top_k=2\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we can create our fusion retriever, which well return the top-2 most similar nodes from the 4 returned nodes from the retrievers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.retrievers import QueryFusionRetriever\n",
    "\n",
    "retriever = QueryFusionRetriever(\n",
    "    [vector_retriever, bm25_retriever],\n",
    "    similarity_top_k=2,\n",
    "    num_queries=4,  # set this to 1 to disable query generation\n",
    "    mode=\"reciprocal_rerank\",\n",
    "    use_async=True,\n",
    "    verbose=True,\n",
    "    # query_gen_prompt=\"...\",  # we could override the query generation prompt here\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply nested async to run in a notebook\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated queries:\n",
      "1. What were the major events or milestones in the history of Interleafe and Viaweb?\n",
      "2. Can you provide a timeline of the key developments and achievements of Interleafe and Viaweb?\n",
      "3. What were the successes and failures of Interleafe and Viaweb as companies?\n"
     ]
    }
   ],
   "source": [
    "nodes_with_scores = retriever.retrieve(\n",
    "    \"What happened at Interleafe and Viaweb?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Score: 0.03 - The UI was horrible, but it proved you could build a whole store through the browser, without any client software or typing anything into the command line on the server.\n",
      "\n",
      "Now we felt like we were really onto something. I had visions of a whole new generation of software working this way. You wouldn't need versions, or ports, or any of that crap. At Interleaf there had been a whole group called Release Engineering that seemed to be at least as big as the group that actually wrote the software. Now you could just update the software right on the server.\n",
      "\n",
      "We started a new company we called Viaweb, after the fact that our software worked via the web, and we got $10,000 in seed funding from Idelle's husband Julian. In return for that and doing the initial legal work and giving us business advice, we gave him 10% of the company. Ten years later this deal became the model for Y Combinator's. We knew founders needed something like this, because we'd needed it ourselves....\n",
      "-----\n",
      "\n",
      "Score: 0.03 - Now we felt like we were really onto something. I had visions of a whole new generation of software working this way. You wouldn't need versions, or ports, or any of that crap. At Interleaf there had been a whole group called Release Engineering that seemed to be at least as big as the group that actually wrote the software. Now you could just update the software right on the server.\n",
      "\n",
      "We started a new company we called Viaweb, after the fact that our software worked via the web, and we got $10,000 in seed funding from Idelle's husband Julian. In return for that and doing the initial legal work and giving us business advice, we gave him 10% of the company. Ten years later this deal became the model for Y Combinator's. We knew founders needed something like this, because we'd needed it ourselves.\n",
      "\n",
      "At this stage I had a negative net worth, because the thousand dollars or so I had in the bank was more than counterbalanced by what I owed the government in taxes. (Had I diligently set aside the proper proportion of the money I'd made consulting for Interleaf?...\n",
      "-----\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for node in nodes_with_scores:\n",
    "    print(f\"Score: {node.score:.2f} - {node.text}...\\n-----\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, both retruned nodes correctly mention Viaweb and Interleaf!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use in a Query Engine!\n",
    "\n",
    "Now, we can plug our retriever into a query engine to synthesize natural language responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.query_engine import RetrieverQueryEngine\n",
    "\n",
    "query_engine = RetrieverQueryEngine.from_args(retriever)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated queries:\n",
      "1. What were the major events or milestones in the history of Interleafe and Viaweb?\n",
      "2. Can you provide a timeline of the key developments and achievements of Interleafe and Viaweb?\n",
      "3. What were the outcomes or impacts of Interleafe and Viaweb on the respective industries they operated in?\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"What happened at Interleafe and Viaweb?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**`Final Response:`** At Interleaf, there was a group called Release Engineering that was as big as the group that actually wrote the software. This suggests that there was a significant focus on managing versions and ports of the software. However, at Viaweb, the founders realized that they could update the software directly on the server, eliminating the need for versions and ports. They started Viaweb, a company that built software that worked via the web. They received $10,000 in seed funding and gave 10% of the company to Julian, who provided the funding and business advice. This deal later became the model for Y Combinator's."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from llama_index.core.response.notebook_utils import display_response\n",
    "\n",
    "display_response(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v3",
   "language": "python",
   "name": "llama_index_v3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
